import wave
import pyaudio
import numpy as np
import matplotlib.pyplot as plt
import pywt

nchannels =0
sampwidth =0
framerate =0
nframes =0

def get_data(name): # 读取读取音频文件信息
    global nchannels
    global sampwidth
    global framerate
    global nframes
    
    
    wave_file = wave.open("./wave_data/"+name+".wav","rb")
    params = wave_file.getparams()
    nchannels, sampwidth, framerate, nframes = params[:4]
    strData = wave_file.readframes(nframes)#读取音频，字符串格式
    waveData = np.fromstring(strData,dtype=np.int16)#将字符串转化为int
    waveData = waveData*1.0/(max(abs(waveData)))#wave幅值归一化
    waveData = np.reshape(waveData,[nframes,nchannels])
    
    return waveData

def draw_time_picture(data,name): # 画出时域图
    time = np.arange(0,len(data))*(1.0 / framerate)
    plt.plot(time,data)
    plt.xlabel("Time(s)")
    plt.ylabel("Amplitude")
    plt.title("Wavedata of "+name)
    plt.grid('on')    
    plt.show()


def draw_fft_picture(data,name): # 画出频域图
    
    nUniquePts = int(np.ceil((len(data)+1)/2.0))
    freqArray = np.arange(0, nUniquePts, 1.0) * (framerate / len(data))
    
    ft=np.fft.fft(data)
    nfft = ft[0:nUniquePts]  
    plt.plot(freqArray, nfft, 'red') 
    plt.xlabel('Freq (Hz)')
    plt.ylabel('Amp. Spectrum') 
    plt.title("FFT1 of "+name)
    plt.grid('on')
    plt.show()
    
    
def draw_fre_picture(fre_data,name,fre_unit=3):
    freqArray = np.arange(0, len(fre_data), 1)*(fre_unit)
    plt.plot(freqArray, fre_data, 'red') 
    plt.xlabel('Freq (Hz)')
    plt.ylabel('Amp. Spectrum') 
    plt.title("Fre_picture of "+name)
    plt.grid('on')
    plt.show()

def filter_index(data,door_amp=0.25,time_diff=3000,move=1000,two_decay_time=16000,start_point=None,end_point=None):  #
    vertex_index = np.where(abs(data[start_point:end_point]) > door_amp)[0]  # 返回所有幅度大于door_amp的时刻
    # print(data)
    result = [] #
    start = vertex_index[0] - move # 得到第一个起点
    end = 0
    for index in range(0, len(vertex_index) - 1):
        if vertex_index[index + 1] - vertex_index[index] > time_diff: # 如果上下两个时刻差大于time_diff，说明他们在是在分别两个小波内
            result.append(start)
            end_pre = vertex_index[index + 1] - move # 把终点往前挪一点，使得其不和小波前端重合
            if end_pre-start > two_decay_time :
                end = start + two_decay_time
            else :
                end = end_pre
            result.append(end)
            start = end_pre + 1
    result.append(start)
    fina_end = len(data)  -10
    if fina_end - start > two_decay_time:
        fina_end = start + two_decay_time
     
    result.append(fina_end)

    return result


def get_fragment(data,two_decay_time=16000, channel=0): # 根据切片的index，切出数据小波
    data = data[:, channel]
    data_vertex_index = filter_index(data) # 得到所有切点
    data_vertex_index = np.array(data_vertex_index)
    

    fragment_result=[]
    for i in range(0,len(data_vertex_index)-1):
        if i%2==0:
            the_fragment = data[ data_vertex_index[i]:data_vertex_index[i+1] ] # 切片数据
            if len(the_fragment)<two_decay_time:
                extend_array = np.zeros(two_decay_time-len(the_fragment))
                fragment_result.append( np.append(the_fragment,  extend_array) )
            else : fragment_result.append(the_fragment)
        
    return fragment_result


def get_fre_info(all_fregement, fre_door=6000, fre_unit=3):
    result = []
    for i in range(0,len(all_fregement)):
        result.append(abs((np.fft.fft(all_fregement[i]))[0:(fre_door//fre_unit)]))
    return result
    


def ave_fre(all_fregement_fre_data, fre_range = 6, fre_unit = 3):
    
    result = []
    for j in range(0, len(all_fregement_fre_data)):
        fre_data = all_fregement_fre_data[j]
        add_data = np.ones(fre_range - (len(fre_data)% fre_range ))
        add_data = add_data*fre_data[-1]
        fre_data = np.append(fre_data, add_data )
        result_ele =np.zeros(len(fre_data)//( fre_range//fre_unit ))
 
        for i in range(0,(fre_range//fre_unit) ):
            result_ele=  result_ele + fre_data[i: :( fre_range//fre_unit )] 
            
        result.append(result_ele/(fre_range//fre_unit))
    return result

    
def save_info(data,name,channel=0):
    data_fragment = np.array(get_fragment(data,channel=channel))
    fre_info = np.array(get_fre_info(data_fragment))
    #eact_fre = ave_fre( fre_info )
    
    for i in range(0,len(data_fragment)):
        np.savetxt("./fre_info/FI_CH"+str(channel)+"_"+name+"_Frg_"+str(i)+".txt",fre_info,fmt="%.3f") # 存进 fre_info 文件夹
        np.savetxt("./tmie_info/TI_CH"+str(channel)+"_"+name+"_Frg_"+str(i)+".txt",data_fragment,fmt="%.3f") # 存进 tmie_info 文件夹
        


